


from langchain_openai import ChatOpenAI
from langchain.schema import SystemMessage
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
# InMemoryVectorStore  基于内存的向量数据库
from langchain_core.vectorstores import InMemoryVectorStore
from langchain.schema.runnable import RunnableMap, RunnableBranch, RunnableLambda
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_huggingface import HuggingFaceEmbeddings
from dotenv import load_dotenv
import os

load_dotenv()


class TravelQASystem:
    def __init__(self, openai_api_key, serpapi_api_key, embed_path):
        """初始化旅游问答系统核心组件"""

        # 初始化语言模型
        self.llm = ChatOpenAI(api_key=openai_api_key,
                              base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
                              model="qwen-plus")

        # 初始化搜索工具
        self.search = TavilySearchResults(tavily_api_key=serpapi_api_key)

        # 初始化嵌入模型
        self.embeddings = HuggingFaceEmbeddings(model_name=embed_path)

        # 构建景点知识库
        self.attraction_data = [
            "故宫：北京地标，明清皇宫，开放时间8:30-17:00",
            "颐和园：皇家园林，昆明湖、长廊等景点",
            "八达岭长城：距离市区70公里，建议游览3-4小时"
        ]

        # 使用内存型向量存储类
        self.vector_store = InMemoryVectorStore.from_texts(
            self.attraction_data, self.embeddings, k=1
        )

    def setup_runnable_pipeline(self):
        """定义Runnable流程管道"""
        # 3.1 问题解析模块：识别地点与查询类型
        parse_prompt = ChatPromptTemplate.from_messages([
            SystemMessage(content="你是旅游助手，需从用户问题中提取地点和查询类型（天气/景点介绍/行程规划）"),
            ("user", """问题：{user_question}请以JSON格式返回：{{"location": "地点", "type": "查询类型"}}""")
        ])
        # {"location": "故宫", "type": "天气"}
        parse_module = parse_prompt | self.llm | JsonOutputParser()  # Output JSON string

        # 3.2 并行获取数据:  天气 + 景点检索
        # 根据大模型的返回的数据   进行天气搜索
        weather_query = RunnableLambda(
            lambda x: self.search.invoke(f"{x['location']} 今日天气")
        )
        # 根据大模型的返回的数据   进行景点检索
        attraction_retrieval = (lambda x: x["location"]) | self.vector_store.as_retriever() | (
            lambda x: x[0].page_content)
        # 创建并行的  Runnable对象
        data_acquisition = RunnableMap({
            "weather": weather_query,
            "attraction": attraction_retrieval,
            "location": (lambda x: x["location"])
        })
        # {"location": "故宫", "attraction": "故宫：北京地标，明清皇宫，开放时间8:30-17:00", "weather": ''空气 良 东风 微风\n\n*    今天 (05-25)   阴  16/28℃ 东南风 微风 \n\n*    明天 (05-26)   阴转多云  19/31℃ 北风 微风 \n\n*    周二 (05-27)   阴转多云  20/33℃ 西风 微风 \n\n*    周三 (05-28)   晴  21/34℃ 南风 微风 \n\n*    周四 (05-29)   阴  23/34℃ 南风 微风 \n\n*    周五 (05-30)   阴  23/33℃ 南风 微风 \n\n*    周六 (05-31)   多云转晴  22/32℃ 南风 微风 \n\n*    周日 (06-01)   多云转阴  23/34℃ 北风 微风 \n\n*    周一 (06-02)   阴转小雨  18/26℃ 北风 微风 \n\n*    周二 (06-03)   小雨转阴  19/27℃ 东南风 微风 \n\n*    周三 (06-04)   阴转晴  20/32℃ 南风 微风 \n\n*    周四 (06-05)   多云转阴  22/34℃ 东南风 微风 [...] *    周五 (06-06)   阴  24/36℃ 南风 微风 \n\n*    周六 (06-07)   阴  24/37℃ 东南风 微风 \n\n*    周日 (06-08)   阴  23/27℃ 北风 微风 \n\n15日天气\n\n[![Image 1]'}

        # 3.3 回答生成模块：整合信息并格式化
        generate_prompt = ChatPromptTemplate.from_messages([
            SystemMessage(content="你是专业旅游顾问，需结合景点信息和天气生成建议"),
            ("user", """地点：{location}
                        景点信息：{attraction}
                        天气情况：{weather}
                        请生成1条行程建议，包含注意事项（如天气相关准备）""")
        ])
        generate_module = generate_prompt | self.llm | (lambda x: x.content.strip())

        # 3.4 串联执行
        # 条件不成立的话     {"location": "故宫", "attraction": "故宫：北京地标，明清皇宫，开放时间8:30-17:00"}
        # 条件成立    {"location": "故宫", "type": "天气"}
        # 条件成立 返回 {"location": "故宫", "attraction": "故宫：北京地标，明清皇宫，开放时间8:30-17:00", "weather": ''空气 良 东风 微风\n\n*    今天 (05-25)   阴  16/28℃ 东南风 微风 \n\n*    明天 (05-26)   阴转多云  19/31℃ 北风 微风 \n\n*    周二 (05-27)   阴转多云  20/33℃ 西风 微风 \n\n*    周三 (05-28)   晴  21/34℃ 南风 微风 \n\n*    周四 (05-29)   阴  23/34℃ 南风 微风 \n\n*    周五 (05-30)   阴  23/33℃ 南风 微风 \n\n*    周六 (05-31)   多云转晴  22/32℃ 南风 微风 \n\n*    周日 (06-01)   多云转阴  23/34℃ 北风 微风 \n\n*    周一 (06-02)   阴转小雨  18/26℃ 北风 微风 \n\n*    周二 (06-03)   小雨转阴  19/27℃ 东南风 微风 \n\n*    周三 (06-04)   阴转晴  20/32℃ 南风 微风 \n\n*    周四 (06-05)   多云转阴  22/34℃ 东南风 微风 [...] *    周五 (06-06)   阴  24/36℃ 南风 微风 \n\n*    周六 (06-07)   阴  24/37℃ 东南风 微风 \n\n*    周日 (06-08)   阴  23/27℃ 北风 微风 \n\n15日天气\n\n[![Image 1]'}
        self.travel_qa_pipeline = (
            # 阶段1：解析问题
                parse_module
                | (lambda x: {"location": x["location"], "type": x["type"]})
                # 阶段2：并行获取数据（仅当查询类型为天气或行程时触发）   RunnableBranch判断
                | RunnableBranch(
            (lambda x: "天气" in x["type"], data_acquisition),
            lambda x: {"location": x["location"], "attraction": attraction_retrieval.invoke(x)}
        )
                # 阶段3：生成回答
                | generate_module
        )


    def process_user_question(self, user_question):
        """处理用户提问并返回回答"""
        input_data = {"user_question": user_question}
        print(input_data)
        # try:
        response = self.travel_qa_pipeline.invoke(input_data)
        return response



if __name__ == '__main__':
    # 需求   根据用户的需求  查询天气和景点的建议
    tqs = TravelQASystem(os.getenv("api_key"), os.getenv("tavily_key"), r'D:\LLM\Local_model\BAAI\bge-large-zh-v1___5')
    tqs.setup_runnable_pipeline()
    ques1 = '故宫的天气怎么样?'
    ans = tqs.process_user_question(ques1)
    print('问题:{},回答:{}'.format(ques1, ans))